package com.pocket.insight.base.sql.jdbc;

import com.mchange.v2.c3p0.DataSources;
import com.mchange.v2.c3p0.PooledDataSource;
import com.pocket.insight.base.constant.SystemConfigConst;
import com.pocket.insight.base.sql.constant.ConnectionConst;
import com.pocket.insight.base.sql.context.JDBCConfigContext;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Predicate;
import java.util.stream.Collectors;

/**
 * 描述：
 *
 * @author shuize
 * @create 2024/11/29
 */
@Slf4j
public class JdbcConnPool {

    private static final HashMap<String, PresetDataSource> QUERY_POOL_MAP = new HashMap<>();
    private static final HashMap<String, PresetDataSource> ETL_POOL_MAP   = new HashMap<>();

    protected static HashMap<String, PresetDataSource> getPurposePool(PoolPurpose poolPurpose) {
        HashMap<String, PresetDataSource> poolMap = null;
        switch (poolPurpose) {
            case QUERY:
                poolMap = QUERY_POOL_MAP;
                break;
            case ETL:
                poolMap = ETL_POOL_MAP;
                break;
        }
        return poolMap;
    }


    /**
     * 根据 jdbcUrl 直接获取数据连接（不走连接池）
     * @param jdbcUrl       数据连接url
     * @param properties    配置参数
     * @param addPresets    预处理sql
     */
    public static Connection directGetConnection(String jdbcUrl, Properties properties, Predicate<Connection> addPresets) throws SQLException {
        Connection connection = (properties == null) ? DriverManager.getConnection(jdbcUrl) : DriverManager.getConnection(jdbcUrl, properties);
        if (connection != null && addPresets != null) {
            try {
                boolean addResult = addPresets.test(connection);
                if (addResult) {
                    log.debug("apply presets for connection hash {}, {}", Integer.valueOf(getConnectionHashCode(connection)), connection.getMetaData().getURL());
                } else {
                    connection.rollback();
                    log.debug("fail to apply presets for connection hash {}, {}", Integer.valueOf(getConnectionHashCode(connection)), connection.getMetaData().getURL());
                }
            } catch (Exception exception) {}
        }
        return connection;
    }

    /**
     * 根据 jdbcUrl 获取数据连接（走连接池）
     *
     * @param poolPurpose   连接池用途
     * @param jdbcUrl       数据连接url
     * @param properties    配置参数
     * @param addPresets    预处理sql
     * @param pingSql       验证连通性sql
     */
    public static Connection getConnection(PoolPurpose poolPurpose, String jdbcUrl, Properties properties, Predicate<Connection> addPresets, String pingSql) throws SQLException {
        if (poolPurpose == PoolPurpose.NONE) {
            return directGetConnection(jdbcUrl, properties, addPresets);
        }
        String connectionKey = buildConnectionUrlKey(jdbcUrl, properties);
        PresetDataSource dataSource = getPurposePool(poolPurpose).get(connectionKey);
        if (dataSource == null) {
            dataSource = getDataSource(poolPurpose, jdbcUrl, properties, pingSql, addPresets);
        }
        long end = 0L;
        long start = System.currentTimeMillis();
        int errorCount = 0;
        try {
            Connection connection = dataSource.getConnection();
            end = System.currentTimeMillis() - start;
            return connection;
        } catch (Exception e) {
            log.error("获取数据连接池连接失败，url【{}】,错误详情：{}", jdbcUrl, e.getMessage());
            end = System.currentTimeMillis() - start;
            errorCount = 1;
            throw e;
        } finally {
            log.error("获取连接耗时： {} ms, 尝试次数： {}", end, errorCount);
        }
    }

    /**
     * 获取连接池
     * @param poolPurpose   连接池用途
     * @param jdbcUrl       数据连接url
     * @param properties    配置参数
     * @param addPresets    预处理sql
     * @param pingSql       验证连通性sql
     */
    public static PresetDataSource getDataSource(PoolPurpose poolPurpose, String jdbcUrl, Properties properties, String pingSql, Predicate<Connection> addPresets) throws SQLException {
        Integer maxConnNum = 10;
        String connectionKey = buildConnectionUrlKey(jdbcUrl, properties);
        if (properties != null) {
            maxConnNum = (Integer) properties.getOrDefault(ConnectionConst.MAX_CONN_NUM, maxConnNum);
        }
        HashMap<String, PresetDataSource> poolMap = getPurposePool(poolPurpose);
        PresetDataSource dataSource = poolMap.get(connectionKey);
        if (dataSource == null) {
            synchronized (poolMap) {
                dataSource = poolMap.get(connectionKey);
                if (dataSource == null) {
                    DataSource dsUnPooled = getUnpooledDataSource(jdbcUrl, properties, addPresets);
                    Map<String, Object> overrides = new HashMap<>();
                    // 连接池初始化时创建的连接数。设置为 0，表示不预先创建连接
                    overrides.put("initialPoolSize", 1);
                    // 连接池最小连接数。设置为 0，表示不限制最小连接数
                    overrides.put("minPoolSize", 1);
                    // 连接池最大连接数。TODO maxPoolSize = 1 时 连接验证特别慢
                    overrides.put("maxPoolSize", maxConnNum);
                    // 每次获取连接增量：连接池中的连接数不足时，每次增加的连接数
                    overrides.put("acquireIncrement", 1);
                    // 最大空闲时间：连接在池中空闲的最大时间（以秒为单位）。超过该时间后，连接将被移除
                    overrides.put("maxIdleTime", 300);
                    // 辅助线程数：用于管理连接池的辅助线程数
                    overrides.put("numHelperThreads", 4);
                    // 空闲连接测试周期：测试空闲连接的时间间隔（以秒为单位）
                    overrides.put("idleConnectionTestPeriod", SystemConfigConst.JDBC_POOL_IDLE_CONNECTION_TEST_PERIOD);
                    // 归还连接时测试：测试连接是否还处于可用状态
                    overrides.put("testConnectionOnCheckin", SystemConfigConst.JDBC_POOL_TEST_CONNECTION_ON_CHECKIN);
                    // 获取连接时测试：测试连接是否还处于可用状态
                    overrides.put("testConnectionOnCheckout", SystemConfigConst.JDBC_POOL_TEST_CONNECTION_ON_CHECKOUT);
                    // 连接获取失败时重试次数
                    overrides.put("acquireRetryAttempts", 2);
                    // 连接获取超时时间
                    overrides.put("checkoutTimeout", SystemConfigConst.JDBC_POOL_CHECKOUT_TIMEOUT);
                    // 首选测试查询语句
                    if (StringUtils.isNotBlank(pingSql) && SystemConfigConst.JDBC_POOL_USE_PING_AS_TEST_SQL.booleanValue()) {
                       // overrides.put("preferredTestQuery", pingSql);
                    }
                    PooledDataSource pooledDataSource = (PooledDataSource) DataSources.pooledDataSource(dsUnPooled, overrides);
                    dataSource = new PresetDataSource(pooledDataSource);
                    log.info("创建数据源连接池，url: 【{}】", jdbcUrl);
                    poolMap.put(connectionKey, dataSource);
                }
            }
        }
        return dataSource;
    }


    /**
     * 创建一个非连接池的数据源对象
     * @param jdbcUrl
     * @param properties
     * @param addPresets
     */
    protected static DataSource getUnpooledDataSource(String jdbcUrl, Properties properties, Predicate<Connection> addPresets) throws SQLException {
        Properties propC3P0 = JDBCConfigContext.removeInnerProperties(properties);
        DataSource unPooledDs = (properties == null) ? DataSources.unpooledDataSource(jdbcUrl) : DataSources.unpooledDataSource(jdbcUrl, propC3P0);
        return new DriverMangerDataSourceAdapter(unPooledDs, addPresets, properties);
    }

    private static int getConnectionHashCode(Connection connection) {
        return connection != null ? connection.hashCode() : 0;
    }

    /**
     * 构建连接池的key（url + 排序后参数：url-key1=value1-key2=value2-...）
     * @param url
     * @param properties
     * @return
     */
    public static String buildConnectionUrlKey(String url, Properties properties) {
        List<String> elements = new ArrayList<>(Arrays.asList(url));
        if (properties != null) {
            elements.addAll(
                    properties.entrySet().stream()
                            .sorted(Comparator.comparing(s -> String.valueOf(s.getKey())))
                            .map(e -> e.getKey() + "=" + e.getValue())
                            .sorted().collect(Collectors.toList())
            );
        }
        return String.join("-", elements);
    }

    public static class PresetDataSource {
        private static final int MAX_KNOWN_CONNECTIONS = 100;
        private final PooledDataSource dataSource;
        private final Set<Integer> concurrentKnownConnections;

        PresetDataSource(PooledDataSource dataSource) {
            new ConcurrentHashMap();
            this.concurrentKnownConnections = ConcurrentHashMap.newKeySet();
            this.dataSource = dataSource;
        }

        Connection getConnection() throws SQLException {
            Connection connection = this.dataSource.getConnection();
            int connectionHashCode = JdbcConnPool.getConnectionHashCode(connection);
            if (connection != null && !this.concurrentKnownConnections.contains(connectionHashCode)) {
                if (this.concurrentKnownConnections.size() > 100) {
                    this.concurrentKnownConnections.clear();
                }
                this.concurrentKnownConnections.add(connectionHashCode);
            }
            return connection;
        }

        void close() throws SQLException {
            this.concurrentKnownConnections.clear();
            this.dataSource.close();
        }
    }

    /**
     * 连接池用途
     */
    public enum PoolPurpose {
        QUERY, ETL, NONE;
    }

}
